/*
 * Copyright 2022 The Open Islands Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.openislands.oi.master.scheduler;

import com.fasterxml.jackson.core.JsonProcessingException;
import org.openislands.oi.config.MyNodeInfo;
import org.openislands.oi.constant.*;
import org.openislands.oi.manager.CommonJobManager;
import org.openislands.oi.scheduling.*;
import org.openislands.oi.dao.entity.TaskDO;
import org.openislands.oi.dao.mapper.TaskMapper;
import org.openislands.oi.error.LogicException;
import org.openislands.oi.pojo.dto.*;
import org.openislands.oi.util.CustomerLogger;
import org.openislands.oi.util.CustomerLoggerFactory;
import org.openislands.oi.util.ObjectUtils;
import org.openislands.oi.util.ValueCopyUtils;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.*;
import java.util.stream.Collectors;

@Component("MasterTaskScheduler")
public class TaskScheduler {
    private static final CustomerLogger schedulingLogger = CustomerLoggerFactory.getLogger(LogTraceType.JOB, LogType.SCHEDULING);
    @Resource
    private CommonJobManager commonJobManager;
    @Resource
    private StateEngine stateEngine;

    @Resource
    private TaskCoordinationProtocol taskCoordinationProtocol;

    @Resource
    private MyNodeInfo myNodeInfo;

    @Resource
    private TaskMapper taskMapper;

    public Boolean createJob(DAG dag, JobDTO jobDTO) {
        jobDTO.setLeaderNode(myNodeInfo.getNodeId());
        Map<RoleType, List<String>> nodes = new HashMap<>();
        dag.getRole().toMap().forEach((roleStr, nodeIds) -> {
            if(nodeIds!=null && !nodeIds.isEmpty()){
                nodes.put(RoleType.valueOf(roleStr.toUpperCase()), nodeIds);
            }
        });
        jobDTO.setNodes(nodes);
        CoordinationResult<Boolean> coordinationResult = taskCoordinationProtocol.createJob(jobDTO);
        if (coordinationResult.isSuccess()) {
            return true;
        } else {
            throw new LogicException(coordinationResult.getCode());
        }
    }

    public Boolean stopJob(JobDTO jobDTO, JobState toState){
        CoordinationResult<Boolean> coordinationResult = taskCoordinationProtocol.stopJob(jobDTO, toState);
        return coordinationResult.isSuccess();
    }

    private void fillCommonResourceParams(JobDTO jobDTO){
        //todo: resource params
    }

    public Boolean deleteJob(JobDTO jobDTO) {
        CoordinationResult<Boolean> coordinationResult = taskCoordinationProtocol.deleteJob(jobDTO);
        return coordinationResult.isSuccess();
    }

    public Boolean updateJob(JobDTO jobDTO) {
        CoordinationResult<Boolean> coordinationResult = taskCoordinationProtocol.syncJob(jobDTO);
        return coordinationResult.isSuccess();
    }

    public Boolean updateJobState(JobDTO jobDTO, JobState jobState){
        CoordinationResult<Boolean> coordinationResult = taskCoordinationProtocol.syncJobState(jobDTO, jobState);
        return coordinationResult.isSuccess();
    }

    public List<JobDTO> queryJob(JobDTO jobDTO) {
        return null;
    }

    public Boolean createTask(JobDTO job, DAG dag, DAGOperator dagOperator) {
        TaskDTO taskDTO = ValueCopyUtils.copyProperties(job, new TaskDTO());
        taskDTO.setTaskName(dagOperator.getName());
        taskDTO.setTaskVersion(0);
        //todo: get from job default conf
        taskDTO.setAutoRetries(3);
        //todo: Operators used by the node should be pre-coordinated
        OperatorDTO operatorDTO = new OperatorDTO(dagOperator.getModule());
        // todo: support secure environment
        taskDTO.setExecutionEnvironment(ExecutionEnvironment.NORMAL);
        taskDTO.setOperatorType(operatorDTO.getType());
        taskDTO.setOperatorAlgorithm(operatorDTO.getAlgorithm());
        //todo: the format depends on whether the operator is dependent on other vendor operators
        if (taskDTO.getOperatorAlgorithm() == AlgorithmType.READER){
            taskDTO.setOutputDataFormat(DataFormat.CSV);
        }else {
            taskDTO.setOutputDataFormat(DataFormat.FATE);
        }
        taskDTO.setOutputModelFormat(ModelFormat.FATE);
        Map<RoleType, List<Map<String, DAGOperator>>> tmp = new HashMap<>();
        // todo: gets participating nodes and their roles from the DAGOperator
        taskDTO.getNodes().forEach((role, nodeIds) -> {
            tmp.put(role, new ArrayList<>());
            nodeIds.forEach(nodeId -> {
                try {
                    Map<String, DAGOperator> t = new HashMap<>();
                    t.put(nodeId, dag.filterOperatorByTriple(dagOperator.getName(), role.name().toLowerCase(), nodeId));
                    tmp.get(role).add(t);
                } catch (JsonProcessingException e) {
                    throw new RuntimeException(e);
                }
            });
        });
        taskDTO.setTaskConfig(ObjectUtils.object2Json(tmp));
        taskDTO.setState(TaskState.READY);
        taskDTO.setNextState(TaskState.READY);
        taskDTO.setCreateTime(new Date());

        CoordinationResult<Boolean> createResult = taskCoordinationProtocol.createTask(taskDTO);
        if (!createResult.isSuccess()) {
            schedulingLogger.error(taskDTO.getJobId(), "create task failed: {}", createResult.getMessage());
            return false;
        }
        return true;
    }

    public Boolean submitTask(JobDTO job, String taskName, Integer taskVersion) {
        //todo: deal task version
        TaskDTO taskDTO = commonJobManager.queryTask(job.getJobId(), taskName, taskVersion);
        TaskState toState = TaskState.SUBMITTED;
        boolean isSuccess;
        //todo: check toState signal timeout
        if (this.actionDrivenByState(taskDTO, toState)){
            CoordinationResult<OfNodeTaskDTO> startResult = taskCoordinationProtocol.startTask(taskDTO, toState);
            if (!startResult.isSuccess()) {
                schedulingLogger.error(job.getJobId(), "start task failed: {}", startResult.getMessage());
                this.actionDrivenByState(taskDTO, TaskState.SUBMISSION_FAILED);
                this.actionDrivenByState(taskDTO, TaskState.UP_FOR_RETRY);
                toState = TaskState.SUBMISSION_FAILED;
                isSuccess = false;
            }else {
                isSuccess = true;
            }
        }else {
            return false;
        }
        CoordinationResult<Boolean> updateStateRes = taskCoordinationProtocol.syncTaskState(taskDTO, toState);
        if(!updateStateRes.isSuccess()){
            schedulingLogger.warn(job.getJobId(), "sync task state error: {}", updateStateRes.getMessage());
        }
        CoordinationResult<Boolean> updateResult = taskCoordinationProtocol.syncTask(taskDTO);
        if (!updateResult.isSuccess()) {
            schedulingLogger.warn(job.getJobId(), "sync task info failed: {}", updateResult.getMessage());
        }
        return isSuccess;
    }

    public TaskDTO stopTask(String taskName, Integer taskVersion, TaskState state) {
        return new TaskDTO();
    }

    public TaskSchedulingResult schedulingTasks(String jobId) {
        //todo: only get latest version
        final List<TaskDTO> tasks = commonJobManager.queryTaskByJob(jobId);
        schedulingLogger.info(jobId, "have {} tasks to schedule", tasks.size());
        TaskSchedulingResult schedulingResult = new TaskSchedulingResult();
        for (TaskDTO task: tasks){
            if (stateEngine.isUnfinishedTaskState(task.getState())){
                final CoordinationResult<OfNodeTaskDTO> collectResult = taskCoordinationProtocol.queryTaskOfNodes(task);
                if (!collectResult.isSuccess()){
                    schedulingLogger.warn(jobId, "collect task {} {} {} failed, next round", task.getJobId(), task.getTaskName(), task.getTaskVersion());
                    continue;
                }
                Set<TaskState> nodeTaskStates = new HashSet<>();
                collectResult.getNodesResult().values().forEach(nodeResults -> {
                    nodeResults.forEach(nodeResult -> {
                        nodeTaskStates.add(nodeResult.getResult().getData().getState());
                    });
                });
                String taskStatesMsg = nodeTaskStates.stream().map(TaskState::name).collect(Collectors.joining(","));
                TaskState currentState = task.getState();
                TaskState newState;
                try{
                    newState = this.calculateTaskState(nodeTaskStates);
                    schedulingLogger.info(jobId, "task current {}, calculate task {} new state by state of all nodes: {}", currentState, newState, taskStatesMsg);
                }catch (Exception e){
                    schedulingLogger.info(jobId, "task current {}, calculate task new state failed by state of all nodes: {}, msg: {}", currentState, taskStatesMsg, e.getMessage());
                    throw e;
                }
                boolean ifStateChange = currentState != newState;
                if (ifStateChange){
                    taskCoordinationProtocol.syncTaskState(task, newState);
                    currentState = newState;
                    task.setState(newState);
                    if (newState == TaskState.SUCCESS){
                        schedulingResult.haveNewSuccess = true;
                    }
                }
                schedulingResult.putState(task.getTaskName(), currentState);
                if (currentState == TaskState.READY){
                    schedulingResult.addReady(task);
                }
                if (ifStateChange && stateEngine.isFinishedTaskState(newState)) {
                    taskCoordinationProtocol.stopTask(task, newState);
                }
            }else {
                schedulingResult.putState(task.getTaskName(), task.getState());
            }
        }
        return schedulingResult;
    }

    private TaskState calculateTaskState(Set<TaskState> nodeTaskStates) {
        if (nodeTaskStates.size() == 1) {
            return nodeTaskStates.iterator().next();
        }
        // have interrupt
        List<TaskState> withInterrupt = stateEngine.intersectionWithTaskStates(nodeTaskStates, stateEngine.getInterruptTaskStates());
        if (!withInterrupt.isEmpty()){
            return withInterrupt.get(0);
        }
        // running/success/not interrupt
        if (nodeTaskStates.contains(TaskState.RUNNING) || nodeTaskStates.contains(TaskState.SUCCESS)){
            return TaskState.RUNNING;
        }
        //different unfinished
        List<TaskState> withUnfinished = stateEngine.intersectionWithTaskStates(nodeTaskStates, stateEngine.getUnfinishedTaskStates());
        if (!withUnfinished.isEmpty()){
            return withUnfinished.get(0);
        }
        throw new LogicException(MessageCodeEnum.DAG_SCHEDULING_ERROR, "no rule to calculate task state");
    }

    private Boolean actionDrivenByState(TaskDTO taskDTO, TaskState nextState){
        CommonResult<Object> result = stateEngine.updateState(taskMapper, nextState, TaskDO::getNextState, TaskState::name, lambdaWrapper -> lambdaWrapper.eq(TaskDO::getJobId, taskDTO.getJobId()).eq(TaskDO::getTaskName, taskDTO.getTaskName()).eq(TaskDO::getTaskVersion, taskDTO.getTaskVersion()));
        if (!result.isSuccess()){
            schedulingLogger.info(taskDTO.getJobId(), "try to get state from {} to {} action token, but get by another scheduler", taskDTO.getState(), nextState);
        }
        return true;
    }

    public static class TaskSchedulingResult{
        private final Map<String, TaskState> taskStates;
        private final List<TaskDTO> readyTasks;
        private boolean haveNewSuccess = false;
        public TaskSchedulingResult(){
            this.taskStates = new HashMap<>();
            this.readyTasks = new ArrayList<>();
        }

        public void putState(String taskName, TaskState state){
            this.taskStates.put(taskName, state);
        }

        public void addReady(TaskDTO taskDTO){
            this.readyTasks.add(taskDTO);
        }

        public Map<String, TaskState> getTaskStates() {
            return taskStates;
        }

        public List<TaskDTO> getReadyTasks() {
            return readyTasks;
        }

        public Set<TaskState> getStates(){
            return new HashSet<>(this.taskStates.values());
        }

        public boolean isHaveNewSuccess() {
            return haveNewSuccess;
        }
    }
}
